package com.lanyun.aicorrectassistant;

import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import io.reactivex.Flowable;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.alibaba.dashscope.aigc.generation.Generation;

import java.io.*;

import java.util.Collections;

public class AssignmentGradingDemo {
    private static final Logger logger = LoggerFactory.getLogger(AssignmentGradingDemo.class);

    private static StringBuilder reasoningContent = new StringBuilder(); // 存储思考过程
    private static StringBuilder finalContent = new StringBuilder();   // 存储最终回复
    private static boolean isFirstPrint = true;

    private static final String API_KEY = "sk-45c1893fd2814c62a95ac1183a1abdaa";
    private static final String MODEL = "deepseek-r1";

    public static void main(String[] args) {
        String localFilePath = "/Users/jqchan/Documents/作业1.docx"; // 替换为本地文件路径

        // 模拟从本地读取文件内容
        String fileContent = readAssignmentContent(localFilePath);
        if (fileContent.isEmpty()) {
            logger.error("无法读取作业文件内容");
            return;
        }

        // 批改作业
        gradeAssignment(fileContent);
        System.exit(0);
    }

    private static void gradeAssignment(String fileContent) {
        // 构建提示语
        String systemPrompt = "你是一位专业的AI批改助手，请对以下学生作业进行批改，并提供详细的批改结果。";
        String fullPrompt = systemPrompt + "\n【学生作业内容】\n" + fileContent + "\n请以自然语言文字形式回复批改结果";

        // 构建消息和参数
        Message userMsg = Message.builder().role(Role.USER.getValue()).content(fullPrompt).build();
        GenerationParam param = buildGenerationParam(userMsg);

        // 创建Generation实例
        Generation gen = new Generation();
        try {
            Flowable<GenerationResult> result = gen.streamCall(param);
            result.blockingForEach(message -> handleGenerationResult(message));
            printFinalResult();
        } catch (NoApiKeyException | ApiException | InputRequiredException e) {
            logger.error("API调用异常：{}", e.getMessage());
        }
    }

    private static GenerationParam buildGenerationParam(Message userMsg) {
        return GenerationParam.builder()
                .apiKey(API_KEY)
                .model(MODEL)
                .messages(Collections.singletonList(userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .build();
    }

    private static void handleGenerationResult(GenerationResult message) {
        String reasoning = message.getOutput().getChoices().get(0).getMessage().getReasoningContent();
        String content = message.getOutput().getChoices().get(0).getMessage().getContent();

        if (!reasoning.isEmpty()) {
            reasoningContent.append(reasoning);
            if (isFirstPrint) {
                System.out.println("====================思考过程====================");
                isFirstPrint = false;
            }
            System.out.print(reasoning);
        }

        if (!content.isEmpty()) {
            finalContent.append(content);
            if (!isFirstPrint) {
                System.out.println("\n====================完整回复====================");
                isFirstPrint = true;
            }
            System.out.print(content);
        }
    }

    private static void printFinalResult() {
        if (finalContent.length() > 0) {
            System.out.println("\n\n=== 最终批改结果 ===");
            System.out.println(finalContent.toString());
        }
    }

    private static String readAssignmentContent(String localFilePath) {
        try (FileInputStream fis = new FileInputStream(localFilePath);
             XWPFDocument document = new XWPFDocument(fis);
             XWPFWordExtractor extractor = new XWPFWordExtractor(document)) {
            return extractor.getText();
        } catch (IOException e) {
            logger.error("读取作业文件失败：{}", e.getMessage());
            return "";
        }
    }
}